import ctypes
import os.path as osp
from enum import Enum
import numpy as np

__all__ = ['Segment', 'ModelsFlag']


class ModelsFlag(Enum):
    SegVesselCatheterBinary_FP16 = 1
    SegDynamicCardiacCatheterBinary_FP16 = 2
    SegPeripheralCatheterBinary_FP16 = 3
    SegCardiacWire_FP16 = 4
    SegVesselCatheterBinary_FP16_Fast = 5
    SegCardiacCatheterAnimal_FP16 = 6
    reserve2 = 7
    reserve3 = 8
    reserve4 = 9
    reserve5 = 10
    length = 11


class SegClassType(ctypes.Structure):
    _fields_ = [
        ('all', ctypes.c_int16),
        ('background', ctypes.c_int16),
        ('foreground', ctypes.c_int16),
        ('air', ctypes.c_int16),
        ('shadowEdge', ctypes.c_int16),
        ('vessel', ctypes.c_int16),
        ('catheter', ctypes.c_int16),
        ('other', ctypes.c_int16),
        ('length', ctypes.c_uint16),
    ]


class Segment(object):
    def __init__(self, model_id: ModelsFlag, logger_name=r"seg"):
        self.dll_handle = ctypes.windll.LoadLibrary(osp.join(osp.dirname(__file__), "dll", "G2DSegment.dll"))
        self.logger_name = logger_name
        self.model_id = model_id.value
        # 获取seg_handle

        # 类型定义
        self.dll_handle.createSegHandle.restype = ctypes.POINTER(ctypes.c_void_p)
        self.dll_handle.createSegHandle.argtypes = [ctypes.c_char_p]

        self.dll_handle.destroySegHandle.argtypes = [ctypes.c_void_p]
        # self.dll_handle.createSegHandle.restype = ctypes.c_int

        self.dll_handle.initialization.argtypes = [ctypes.c_void_p, ctypes.c_int32, ctypes.c_uint16, ctypes.c_uint16]
        self.dll_handle.initialization.restype = ctypes.c_int

        self.dll_handle.getSegClass.restype = SegClassType  # 传递segment::G2DSegmentHandler*类型的参数
        self.dll_handle.getSegClass.argtypes = [ctypes.c_void_p]

        self.dll_handle.getThreshSegImg.argtypes = [
            ctypes.c_void_p,
            np.ctypeslib.ndpointer(dtype=np.uint16),
            np.ctypeslib.ndpointer(dtype=np.uint16),
            ctypes.c_int
        ]
        self.dll_handle.getThreshSegImg.restype = ctypes.c_int

        self.dll_handle.reloadModel.restype = ctypes.c_int
        self.dll_handle.reloadModel.argtypes = [ctypes.c_void_p, ctypes.c_int32]

        # 初始化的相关操作
        self.seg_handle = self.__getSegHandle()                    # 获取分割句柄
        self.initialization(1024, 1024)               # 加载模型

    def __getSegHandle(self):
        return self.dll_handle.createSegHandle(self.logger_name.encode())

    def __destroySegHandle(self):
        self.dll_handle.destroySegHandle(self.seg_handle)

    def initialization(self, height, width):
        self.dll_handle.initialization(self.seg_handle, self.model_id, height, width)

    def getSegClass(self):
        result = self.dll_handle.getSegClass(self.seg_handle)
        return result

    def process(self, img: np.ndarray, ioi):
        if img.dtype != np.uint16:
            raise TypeError(f"仅支持uint16数组输入，如若非uint16请将像素rescale至0-65535并转为uint16类型进行输入")
        height, width = img.shape
        self.initialization(height, width)          # 初始化当前尺寸
        result_img = np.zeros_like(img)             # 申请返回数组
        self.dll_handle.getThreshSegImg(self.seg_handle, img, result_img, ioi)
        return result_img

    def reloadModel(self,  model_id: ModelsFlag):
        self.model_id = model_id.value
        self.dll_handle.reloadModel(self.seg_handle, self.model_id)

    def __del__(self):
        self.__destroySegHandle()





